#perform MK tests

#see get_transcripts.txt for how this text file was created
chick.transcripts <- read.table('/Paterson/Datafiles/grouse/seq_cap/blast_out/chick_prot_cDNA.txt',header=TRUE, stringsAsFactors = FALSE)

#see cds_translate for some of the function such as fixedSubs()

grouse.chick.map <- merge(grouse600annot[,c("grouse.id","ensembl.id")],chick.transcripts,by.x="ensembl.id",by.y="ensembl_peptide_id")

#get chicken transcripts
chick.seqs <- read.BStringSet("/Paterson/Datafiles/grouse/seq_cap/blast_out/chick_transcripts.fasta")

grouse.chick.MK <- grouse.chick.map
grouse.chick.MK[,c("Dn", "Ds", "transitions", "transversions", "dN", "dN.se", "dS", "dS.se", "kappa", "omega", "Ls", "Ln", "t")] <- -1

#useful code to set PATH variable
Sys.setenv(PATH=paste(Sys.getenv("PATH"),"/usr/local/bioinf/paml/bin",sep=":"))


for(grouse.gene in 1:nrow(grouse.chick.map)){#change to nrow(grouse.chick.map)
	chick.transcript <- as.character(chick.seqs[[grep(grouse.chick.map$ensembl_transcript_id[grouse.gene],names(chick.seqs))]])
	grouse.transcript <- getCDS(seq.name=grouse.chick.map$grouse.id[grouse.gene],single.cds=TRUE)
	if(!"transcript" %in% names(grouse.transcript)) next
	
	#remove *s from protein/transcript
	star.pos <- grep('\\*',strsplit(grouse.transcript$protein,'')[[1]])
	if(length(star.pos)>0){
		tst.trans <- strsplit(grouse.transcript$transcript[[1]],'')[[1]]
		cut.trans <- tst.trans[-c( (star.pos-1)*3+1, (star.pos-1)*3+2, (star.pos-1)*3+3 )]
		grouse.transcript$transcript <- paste(cut.trans,collapse="")
		}
	
	test.MK <- try(fixedSubs(chick.transcript,grouse.transcript$transcript,transition.ratio=TRUE))
	if(class(test.MK)=="try-error") next
	grouse.chick.MK[grouse.gene,c("Dn", "Ds", "transitions", "transversions", "dN", "dN.se", "dS", "dS.se", "kappa", "omega", "Ls", "Ln", "t")] <- test.MK
	}
## NB, this still doesn't work well where there are lots of CDSs predicted by Exonerate
# think about trying again with just the first cds, or remove these from analysis

#also think about removing sequences with <90% (say) homology
#tranalign may not be working perfectly, especially at the end of sequences where
#frameshifts may not be noticed
# exclude sets of sites where >50% of aa are different for sliding window?
# in order to exclude blocks with low homology
# and check whether alignments are 'worse' in IF or non-IF genes

#try again but output separate files that can then be edited
for(grouse.gene in 1:nrow(grouse.chick.map)){#change to nrow(grouse.chick.map)
	chick.transcript <- as.character(chick.seqs[[grep(grouse.chick.map$ensembl_transcript_id[grouse.gene],names(chick.seqs))]])
	grouse.transcript <- getCDS(seq.name=grouse.chick.map$grouse.id[grouse.gene],single.cds=TRUE)
	if(!"transcript" %in% names(grouse.transcript)) next
	
	#remove *s from protein/transcript
	star.pos <- grep('\\*',strsplit(grouse.transcript$protein,'')[[1]])
	if(length(star.pos)>0){
		tst.trans <- strsplit(grouse.transcript$transcript[[1]],'')[[1]]
		cut.trans <- tst.trans[-c( (star.pos-1)*3+1, (star.pos-1)*3+2, (star.pos-1)*3+3 )]
		grouse.transcript$transcript <- paste(cut.trans,collapse="")
		}
	
	test.MK <- try(fixedSubs(chick.transcript,grouse.transcript$transcript,transition.ratio=TRUE))
	if(class(test.MK)=="try-error") next
	tmp.cmd1 <- paste("mv /Paterson/Datafiles/grouse/test_files/transcripts_tmpfile.afa /Paterson/Datafiles/grouse/alignments/",paste(grouse.chick.map[grouse.gene,'grouse.id'],"_transcript.afa",sep=""),sep="")
	tmp.cmd2 <- paste("mv /Paterson/Datafiles/grouse/test_files/proteins_tmpfile.afa /Paterson/Datafiles/grouse/alignments/",paste(grouse.chick.map[grouse.gene,'grouse.id'],"_protein.afa",sep=""),sep="")
	system(tmp.cmd1)
	system(tmp.cmd2)
	grouse.chick.MK[grouse.gene,c("Dn", "Ds", "transitions", "transversions", "dN", "dN.se", "dS", "dS.se", "kappa", "omega", "Ls", "Ln", "t")] <- test.MK
	}



#calculate polymorphims
grouse.chick.MK$Pn <- -1
grouse.chick.MK$Ps <- -1
grouse.chick.MK$fisher.pv <- -1
grouse.chick.MK$fisher.or <- -1
grouse.chick.MK$fisher.lower <- -1
grouse.chick.MK$fisher.upper <- -1

for(mki in 1:nrow(grouse.chick.MK)){
	if(!grouse.chick.MK$grouse.id[mki] %in% names(grouseSNPdfs)) next
	tmp.snp <- QCsnp(grouseSNPdfs[[grouse.chick.MK$grouse.id[mki]]])
	if(nrow(tmp.snp)==0){
		grouse.chick.MK[mki,c("Pn","Ps")] <- c(0,0)		}else{
		tmp.cnt <- colSums(tmp.snp[,c("cds","dn")])
		grouse.chick.MK[mki,c("Pn","Ps")] <- c(tmp.cnt['dn'],tmp.cnt['cds']-tmp.cnt['dn'])
		tmp.fisher <- try(fisher.test(matrix(as.numeric(grouse.chick.MK[mki,4:7]),nrow=2,byrow=TRUE)))
		if(class(tmp.fisher)=="try-error"){
			cat(mki,"  ",grouse.chick.MK$grouse.id[mki],'  failed')
			next
			}
		grouse.chick.MK$fisher.pv[mki] <- tmp.fisher$p.value
		grouse.chick.MK$fisher.or[mki] <- tmp.fisher$estimate
		grouse.chick.MK$fisher.lower[mki] <- tmp.fisher$conf.int[1]
		grouse.chick.MK$fisher.upper[mki] <- tmp.fisher$conf.int[2]
		}
	
	}




#need to decide what the QC for a coding polymorphism is
# exclude frameshifts
# require >= 30x coverage across all populations

#NB there's also another function called QCsnp in one of the files!

QCsnp <- function(snp.df,coverage = 30,min.freq=0.05,coding=TRUE,exclude.frameshift=TRUE, exclude.duplicate=TRUE){
	#snp.df is produced by codingSNPref
	#coverage is the minumum tsv coverage, summed across all populations, required
	#min.freq is the minium freq needed in any of the populations
	#coding, only limit to coding snps or not
	#frmaeshift, exclude frameshifts as probably sequencing errors
	
	snp.out <- snp.df
	#take out duplicated snps
	if(exclude.duplicate){
		snp.out <- snp.out[snp.out$start.pos%in%as.numeric(names(table(snp.out$start.pos))[table(snp.out$start.pos)==1]),]
		}
	
	#trim out very low freq snps
	snp.out <- snp.out[rowSums(snp.out[,grep("freq",names(snp.out))]>min.freq)>0,]
	
	
	if(coding) snp.out <- snp.out[snp.out$cds,] #cds is T/F
	if(exclude.frameshift){
		tmp.ref <- sapply(X=snp.out$id,FUN=function(X){
			tmp.id <- strsplit(X,"_")[[1]]
			tmp.id[length(tmp.id)-1]
			})
		tmp.var <- sapply(X=snp.out$id,FUN=function(X){
			tmp.id <- strsplit(X,"_")[[1]]
			tmp.id[length(tmp.id)]
			})
		snp.out <- snp.out[tmp.ref!="-"&tmp.var!="-",]
		}
	snp.out
	}

#add Fst data
 Fst.calc <- function(p,sample.size){
	#optionally giive a vector (T/F or 1/0) giving the resistant pops
	#p is a vector of allele frequencies
	
	#do the Fst calculations (see Weir p. 147)
	if(length(sample.size)==1) sample.size <- rep(sample.size,length(sampled.mat))
	n <- sample.size
	r <- length(sample.size)
	p.hat <- sum(n*p)/sum(n)
	n.bar <- sum(n)/r
	
	sa <- sum(n*(p-p.hat)^2)/(r*n.bar) #note this is different from Weir to stop Fst>1
	
	theta <- sa/(p.hat*(1-p.hat))
	theta
	}

grouseFst.df <- vector('list',length(grouseSNPdfs))
names(grouseFst.df) <- names(grouseSNPdfs)
for(grouse.gene in 1:length(grouseSNPdfs)){
	tst.df <- QCsnp(grouseSNPdfs[[grouse.gene]])
	if(nrow(tst.df)==0) next #jump out if no decent snps
	tmp.out <- data.frame(Fst=rep(0,nrow(tst.df)))
	tmp.out <- cbind(tst.df[,c('cDNA','id','start.pos')],tst.df[,grep('freq',names(tst.df))],tst.df[,grep('tsv.depth',names(tst.df))],tmp.out)
	for(snpi in 1:nrow(tmp.out)){
		tmp.out[snpi,'Fst'] <- Fst.calc(tst.df[snpi,grep('freq',names(tst.df))],tst.df[snpi,grep('tsv.depth',names(tst.df))])
		}
	grouseFst.df[[grouse.gene]] <- tmp.out
	}

#get immune hits
grouse.chick.MK$type <- factor('control',levels=c('control','immune'))
grouse.imm.ids <- read.table('/Paterson/Datafiles/grouse/seq_cap/commands/R_code/RG_FINALCLONE_ID_sp4.txt',stringsAsFactors=FALSE,row.names=NULL,col.names=c("id","length"))

grouse.chick.MK$type[grouse.chick.MK$grouse.id %in% grouse.imm.ids$id] <- 'immune'

#add weighted mean Fst for each gene
grouse.chick.MK$Fst.mean <- -1
for(grouse.gene in 1:length(grouseFst.df)){
	if(is.null(grouseFst.df[[grouse.gene]])) next
	depth.weights <- rowSums(grouseFst.df[[grouse.gene]][,grep('depth',names(grouseFst.df[[grouse.gene]]))])
	depth.weights <- depth.weights/sum(depth.weights)
	grouse.chick.MK$Fst.mean[grouse.chick.MK$grouse.id==grouseFst.df[[grouse.gene]]$cDNA[1]] <- sum(depth.weights*grouseFst.df[[grouse.gene]]$Fst)
	}
plot(Fst.mean~type,data=grouse.chick.MK,subset=grouse.chick.MK$Fst.mean> -0.5)

## added 4/8/10
#get data ready for Welch tests
#need a measure of mutational opportunity
#work out ts/tv ratio from alignment between grouse and chicken
#then work along gene counting number of syn and non-syn changes
#see fixed_subs2.R

##added 31/8/10
#produce a trimmed version of alignments
